use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc::SyncSender;

use glam::{dvec3, DVec3};
use rayon::prelude::*;

use crate::{SCREEN_W, SCREEN_H, Hitables, Hitable, Ray, Interval, random_range, linear_to_gamma, random_uint_disk, degrees_to_radians};

/// the builder of Renderer
/// 
#[derive(Debug)]
pub struct RendererBuilder {
    tx: SyncSender<(usize, DVec3)>,

    width: usize,
    height: usize,

    sample_count: usize,
    max_depth: usize,

    camera: DVec3,
    lookat: DVec3,
    up: DVec3,
    fov: f64,

    defocus_angle: f64,
    focus_dist: f64,

    background: DVec3,
}

impl RendererBuilder {
    pub fn new(tx: SyncSender<(usize, DVec3)>) -> Self {
        Self {
            tx,
            width: SCREEN_W,
            height: SCREEN_H,
            sample_count: 10,
            max_depth: 10,
            camera: dvec3(0.0, 0.0, -10.0),
            lookat: dvec3(0.0, 0.0, 0.0),
            up: dvec3(0.0, 1.0, 0.0),
            fov: 90.0,
            defocus_angle: 0.0,
            focus_dist: 10.0,
            background: dvec3(0.7, 0.8, 1.0),
        }
    }

    #[allow(unused)]
    pub fn screen_size(mut self, width: usize, height: usize) -> Self {
        self.width = width;
        self.height = height;
        self
    }

    #[allow(unused)]
    pub fn sample_count(mut self, sample_count: usize) -> Self {
        self.sample_count = sample_count;
        self
    }

    #[allow(unused)]
    pub fn max_depth(mut self, max_depth: usize) -> Self {
        self.max_depth = max_depth;
        self
    }

    #[allow(unused)]
    pub fn camera(mut self, camera: DVec3, lookat: DVec3, up: DVec3, fov: f64) -> Self {
        self.camera = camera;
        self.lookat = lookat;
        self.up = up;
        self.fov = fov;
        self
    }

    #[allow(unused)]
    pub fn dof(mut self, defocus_angle: f64, focus_dist: f64) -> Self {
        self.defocus_angle = defocus_angle;
        self.focus_dist = focus_dist;
        self
    }

    #[allow(unused)]
    pub fn background(mut self, background: DVec3) -> Self {
        self.background = background;
        self
    }

    pub fn build(self) -> Renderer {
        Renderer::new(
            self.width,
            self.height,
            self.camera,
            self.lookat,
            self.up,
            self.fov,
            self.sample_count,
            self.max_depth,
            self.defocus_angle,
            self.focus_dist,
            self.background,
            self.tx,
        )
    }
}

pub struct Renderer {
    tx: SyncSender<(usize, DVec3)>,

    width: usize,
    num_pixel: usize,

    sample_count: usize,
    max_depth: usize,

    defocus_angle: f64,

    center: DVec3,
    pixel_delta_u: DVec3,
    pixel_delta_v: DVec3,
    pixel00_loc: DVec3,
    defocus_disk_u: DVec3,
    defocus_disk_v: DVec3,

    background: DVec3,
}

impl Renderer {
    pub fn new(
        width: usize,
        height: usize,
        camera: DVec3,
        lookat: DVec3,
        up: DVec3,
        fov: f64,
        sample_count: usize,
        depth: usize,
        defocus_angle: f64,
        focus_dist: f64,
        background: DVec3,
        tx: SyncSender<(usize, DVec3)>,
    ) -> Self
    {
        // total number of pixel need ray trace
        let num_pixel = width * height;

        let center = camera;

        let f64_w = width as f64;
        let f64_h: f64 = height as f64;

        // tan(theta / 2) = (viewport_h * 2) / focal_lenght, this for calc viewport_h
        let h = (fov * 0.5).to_radians().tan();

        // frustum plane at focal length
        let viewport_h = 2.0 * h * focus_dist;
        let viewport_w = viewport_h * f64_w / f64_h;

        // base vector of frustum plane in world space
        // note that w is not camera direction, but just one place use it and sub it
        let w = (camera - lookat).normalize();
        let u = up.cross(w).normalize();
        let v = w.cross(u);

        // viewport size base on wuv vector
        let viewport_u = viewport_w * u;
        let viewport_v = viewport_h * -v;

        // one pixel delta base on wuv vector
        let pixel_delta_u = viewport_u / f64_w;
        let pixel_delta_v = viewport_v / f64_h;

        // note that `center - w * focal_length`, use sub operation there
        let viewport_ul = center - w * focus_dist  - viewport_u * 0.5 - viewport_v * 0.5;
        // the location of pixel (0, 0) in world space
        let pixel00_loc = viewport_ul + 0.5 * (pixel_delta_u + pixel_delta_v);

        let defocus_radius = focus_dist * degrees_to_radians(defocus_angle * 0.5).tan();
        let defocus_disk_u = u * defocus_radius;
        let defocus_disk_v = v * defocus_radius;

        Self {
            tx,
            width,
            num_pixel,
            sample_count,
            max_depth: depth,
            defocus_angle,
            center,
            pixel_delta_u,
            pixel_delta_v,
            pixel00_loc,
            defocus_disk_u,
            defocus_disk_v,
            background,
        }
    }

    pub fn render(
        &self,
        hitables: Hitables,
        exit: Arc<AtomicBool>,
    )
    {
        let tx = self.tx.clone();

        (0..self.num_pixel)
            .collect::<Vec<_>>()
            .into_par_iter()
            .take_any_while(|_| !exit.load(Ordering::Relaxed))
            .for_each_with(tx, |tx,idx| {
                let x = (idx % self.width) as f64;
                let y = (idx / self.width) as f64;

                // multi sample
                // tracing multi times per pixel, and average the result
                let mut color = DVec3::ZERO;
                for _ in 0..self.sample_count {
                    color += self.ray_color(&hitables, self.random_ray(x, y), self.max_depth);
                }
                color /= self.sample_count as f64;

                // convert linear color to gamma color, all calculation in linear space before
                linear_to_gamma(&mut color);

                tx.send((idx, color))
                    .expect("send data back to main thread failed");
            });
    }

    /// random ray for pixel in (x, y)
    /// 
    fn random_ray(&self, x: f64, y: f64) -> Ray {
        let offset = Self::sample_square();

        let pixel_sample = self.pixel00_loc 
            + (x + offset.x) * self.pixel_delta_u
            + (y + offset.y) * self.pixel_delta_v;

        let ray_origin = if self.defocus_angle < 0.0 { self.center } else { self.defocus_disk_sample() };
        let ray_dir = pixel_sample - ray_origin;

        Ray::new(ray_origin, ray_dir)
    }

    fn defocus_disk_sample(&self) -> DVec3 {
        let p = random_uint_disk();
        self.center + p.x * self.defocus_disk_u + p.y * self.defocus_disk_v
    }

    /// sample a point in a square
    ///
    #[inline]
    fn sample_square() -> DVec3 {
        let u = random_range(-0.5, 0.5);
        let v = random_range(-0.5, 0.5);
        DVec3::new(u, v, 0.0)
    }

    fn ray_color(&self, scene: &Hitables, ray: Ray, depth: usize) -> DVec3 {
        // just return black if max depth reached
        if depth <= 0 { return DVec3::ZERO }

        scene.hit(&ray, Interval::new(0.001, f64::MAX))
            .and_then(|hit| {
                let emission = hit.mat.emitted(&ray, &hit);

                hit.mat.scatter(&ray, &hit)
                    .and_then(|(attenuation, scattered)| {
                        let scatter = attenuation * self.ray_color(scene, scattered, depth - 1);
                        Some(scatter + emission)
                    })
                    .or(Some(emission))
            })
            .unwrap_or(self.background)
    }
}
